import pandas as pd
import sqlite3
from sqlite3 import Error

"""1. 创建一个示例 DataFrame"""

data = {
    "id": [1, 2, 3, 4, 5],
    "name": ["Alice", "Bob", "Charlie", "David", "Eve"],
    "age": [25, 30, 35, 40, 45],
    "department": ["HR", "Engineering", "Marketing", "Engineering", "HR"],
    "salary": [50000, 75000, 60000, 80000, 55000],
}

df = pd.DataFrame(data)

print("原始 DataFrame:")
print(df)
print("\n")

# 2. 连接到 SQLite 数据库 (如果不存在会自动创建)


def create_connection(db_file):
    """创建数据库连接"""
    conn = None
    try:
        conn = sqlite3.connect(db_file)
        print(f"成功连接到 SQLite 数据库: {db_file}")
        print(f"SQLite 版本: {sqlite3.version}")
        return conn
    except Error as e:
        print(e)

    return conn


database = "company.db"
conn = create_connection(database)

# 3. 将 DataFrame 写入 SQLite 数据库

if conn is not None:
    try:

        # 使用 to_sql 方法将 DataFrame 写入数据库

        # if_exists 参数:

        #   'fail' - 如果表存在则不执行任何操作

        #   'replace' - 如果表存在则删除后重建

        #   'append' - 如果表存在则将数据追加到表中

        df.to_sql("employees", conn, if_exists="replace", index=False)
        print("DataFrame 已成功写入 SQLite 数据库表 'employees'")

        # 验证数据是否已写入
        query = "SELECT * FROM employees"
        df_from_db = pd.read_sql_query(query, conn)

        print("\n从数据库读取的数据:")
        print(df_from_db)

    except Error as e:
        print(e)
    finally:
        # 关闭连接
        conn.close()
        print("\n数据库连接已关闭")

else:
    print("无法创建数据库连接")


print(df[df["salary"] > 60000])
result = df.query("age > 30")
print(result)
